from lightning.pytorch.loggers import WandbLogger, TensorBoardLogger, CSVLogger
from .utils import registry
from typing import Optional


@registry.loggers('wandb')
def build_wandb_logger(name: str, project: str = "omgpt_logs", save_dir: str = './logs'):
    return WandbLogger(name=name, project=project, save_dir=save_dir)


@registry.loggers('csv')
def build_csv_logger(save_dir: str, name: str = 'omgpt_logs'):
    return CSVLogger(save_dir=save_dir, name=name)


@registry.loggers('tensorboard')
def build_tensorboard_logger(save_dir: str, name: Optional[str] = None):
    return TensorBoardLogger(save_dir=save_dir, name=name)